#!/usr/bin/env python3
"""
ProFTPD MCP Server

面向 proftpd FTP 服务器软件的命令行操作封装。

MCP tools 列表：
1. proftpd_service_management - 管理 proftpd 服务的启动、重载和状态检查
   参数：action (str) - 要执行的操作，可选值：start, reload, status
   返回：统一 JSON 结构，包含命令执行结果

2. proftpd_socket_management - 管理 proftpd.socket 的重启和重载操作
   参数：action (str) - 要执行的操作，可选值：restart, reload
   返回：统一 JSON 结构，包含命令执行结果
"""

import subprocess
import json
from typing import Dict, Any
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("ProFTPD MCP Server")

@mcp.tool()
def proftpd_service_management(action: str) -> Dict[str, Any]:
    """
    管理 proftpd 服务的启动、重载和状态检查
    
    参数:
        action: 要执行的操作，可选值：
                "start" - 启动 proftpd 服务
                "reload" - 重载 proftpd 服务配置
                "status" - 检查 proftpd 服务状态
    
    返回 JSON 结构:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 命令标准输出
            "stderr": str         # 命令标准错误
        }
    """
    try:
        if action == "start":
            # 启动 proftpd 服务
            result = subprocess.run(
                ["systemctl", "start", "proftpd.service"],
                capture_output=True,
                text=True,
                timeout=30
            )
        elif action == "reload":
            # 修改服务配置并重载
            subprocess.run(["sed", "-i", "s|ExecStart = /usr/sbin/proftpd|ExecStart = /usr/sbin/proftpd -q|g", 
                          "/usr/lib/systemd/system/proftpd.service"], 
                         capture_output=True, text=True)
            subprocess.run(["systemctl", "daemon-reload"], 
                         capture_output=True, text=True)
            result = subprocess.run(
                ["systemctl", "reload", "proftpd.service"],
                capture_output=True,
                text=True,
                timeout=30
            )
        elif action == "status":
            # 检查服务状态
            result = subprocess.run(
                ["systemctl", "status", "proftpd.service"],
                capture_output=True,
                text=True,
                timeout=30
            )
        else:
            return {
                "success": False,
                "command": f"proftpd_service_management {action}",
                "exit_code": -1,
                "stdout": "",
                "stderr": f"无效的操作: {action}，支持的操作: start, reload, status"
            }
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(result.args),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
        
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": f"systemctl {action} proftpd.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": "命令执行超时"
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"systemctl {action} proftpd.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"命令执行异常: {str(e)}"
        }

@mcp.tool()
def proftpd_socket_management(action: str) -> Dict[str, Any]:
    """
    管理 proftpd.socket 的重启和重载操作
    
    参数:
        action: 要执行的操作，可选值：
                "restart" - 重启 proftpd.socket
                "reload" - 重载 proftpd.socket 配置
    
    返回 JSON 结构:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 命令标准输出
            "stderr": str         # 命令标准错误
        }
    """
    try:
        if action == "restart":
            result = subprocess.run(
                ["systemctl", "restart", "proftpd.socket"],
                capture_output=True,
                text=True,
                timeout=30
            )
        elif action == "reload":
            result = subprocess.run(
                ["systemctl", "reload", "proftpd.socket"],
                capture_output=True,
                text=True,
                timeout=30
            )
        else:
            return {
                "success": False,
                "command": f"proftpd_socket_management {action}",
                "exit_code": -1,
                "stdout": "",
                "stderr": f"无效的操作: {action}，支持的操作: restart, reload"
            }
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(result.args),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
        
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": f"systemctl {action} proftpd.socket",
            "exit_code": -1,
            "stdout": "",
            "stderr": "命令执行超时"
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"systemctl {action} proftpd.socket",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"命令执行异常: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()